# -*- coding:utf-8 -*-
import sys

from loguru import logger
import argparse
import selectors
import socket
import pyDes

logger.add('runtime.log', format="{time:YYYY-MM-DD HH:mm:ss} | {message}")


class Des:
    def __init__(self, key=b'test1234', iv=b"\x00\x00\x00\x00\x00\x00\x00\x00"):
        self.key = key
        self.iv = iv

    def encrypt(self, data):
        k = pyDes.des(self.key, pyDes.CBC, self.iv, pad=None, padmode=pyDes.PAD_PKCS5)
        encrypt_data = k.encrypt(data)
        return encrypt_data

    def decrypt(self, data):
        k = pyDes.des(self.key, pyDes.CBC, self.iv, pad=None, padmode=pyDes.PAD_PKCS5)
        decrypt_data = k.decrypt(data)
        return decrypt_data


class TcpProxy:
    """
    TCP代理程序
    """

    def __init__(self, local_addr, local_port, remote_addr, remote_port, up, down):
        """
        代理初始化
        :param local_addr: 本地监听地址
        :param local_port: 本地监听端口
        :param remote_addr: 远程服务器地址
        :param remote_port: 远程服务器端口
        """
        self.local_addr = local_addr
        self.local_port = local_port
        self.remote_addr = remote_addr
        self.remote_port = remote_port
        self.up = up
        self.down = down

        # 创建本地sock对象 并监听端口，接受客户端请求
        local_sock = socket.socket()
        selector = selectors.DefaultSelector()

        try:
            # 绑定端口，设置监听
            local_sock.bind((local_addr, local_port))
            local_sock.listen(100)
            # 设置为非阻塞模式
            local_sock.setblocking(False)
            self.local_sock = local_sock
            print(f"start listening {local_addr}:{local_port}")
            selector.register(self.local_sock, selectors.EVENT_READ, self.accept)

        except OSError as e:
            logger.error(f"{local_addr}:{local_port} 当前端口被占用，请结束其它占用程序")

        remote_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            remote_sock.connect((remote_addr, remote_port))   # remote-sock is blocking now.
        except socket.error as e:
            print(f"{remote_addr}:{remote_port} 远程连接失败，请检测")
            remote_sock.close()
            exit()
        self.remote_sock = remote_sock
        self.selector = selector
        remote_sock.setblocking(False)

    @staticmethod
    def receive_from(conn):
        """
        返回数据
        :param conn:
        :return:
        """
        buffer = bytes()
        # 设置一个超时时间，如果超过2s没有接收到消息则进入异常处理（必须设置,如果不设置超时时间，则会阻塞，直到接收到消息为止，在本例中也许不会出现问题，但是如果客户端也需要发送请求时就会出现问题）
        # conn.settimeout(0.5)
        try:
            # 将收到的消息拼接成buffer
            while True:
                data = conn.recv(4096)
                if not data:
                    break
                buffer += data
        # 如果2s都没收到消息，就返回空的buffer
        except:
            pass
        return buffer

    def process_client_recv(self, data):
        """处理从Client接收到的数据"""
        if data:
            if self.up == 'Decrypt':
                print(f"当前Client接收到未解密数据：{data}")
                d = Des()
                decrypt_data = d.decrypt(data)
                print(f"当前Client接收到解密后数据：{decrypt_data}")
                return bytes(decrypt_data)
            elif self.up == 'Encrypt':
                print(f"当前Client接收到未加密数据：{data}")
                d = Des()
                encrypt_data = d.encrypt(data)
                print(f"当前Client接收到加密后数据：{encrypt_data}")
                return bytes(encrypt_data)
        return data

    def process_remote_recv(self, data):
        if data:
            if self.down == 'Decrypt':
                print(f"当前Remote接收到未解密数据：{data}")
                d = Des()
                decrypt_data = d.decrypt(data)
                print(f"当前Remote接收到解密后数据：{decrypt_data}")
                return bytes(decrypt_data)
            elif self.down == 'Encrypt':
                print(f"当前Remote接收到未加密数据：{data}")
                d = Des()
                encrypt_data = d.encrypt(data)
                print(f"当前Remote接收到加密后数据：{encrypt_data}")
                return bytes(encrypt_data)
        return data

    def accept(self, sock):
        # 建立客户端连接
        conn, addr = sock.accept()
        print('accepted Client', addr)
        # 设置非阻塞模式
        conn.setblocking(False)
        # 再次注册一个连接，将其加入监测列表中，
        self.selector.register(conn, selectors.EVENT_READ, self.read)

    def read(self, local_conn):
        try:
            # 抛出客户端强制关闭的异常（如手动关闭客户端黑窗口）
            data = self.receive_from(local_conn)
            if data:
                data = self.process_client_recv(data)
                self.remote_sock.send(data)
                receive_data = self.receive_from(self.remote_sock)
                receive_data = self.process_remote_recv(receive_data)
                print(f"receive_data: {receive_data}")
                # 返回远程数据到客户端
                local_conn.send(receive_data)
            else:
                print('Client closed.', local_conn)
                # 将conn从监测列表删除
                self.selector.unregister(local_conn)
                local_conn.close()
        except ConnectionResetError:
            print('Client forcibly closed.', local_conn)
            # 将conn从监测列表删除
            self.selector.unregister(local_conn)
            local_conn.close()

    def run(self):
        while True:
            events = self.selector.select()  # 监测有无活动对象，没有就阻塞在这里等待
            for key, mask in events:  # 有活动对象了
                callback = key.data  # key.data 是注册时传递的 accept 函数
                callback(key.fileobj)  # key.fileobj 就是传递的 socket 对象


if __name__ == '__main__':
    print("TcpProxy.py -h 可查看参数说明\n运行示例: TcpProxy.py -ld 127.0.0.1 -lp 5000 -rd 192.168.1.1 -rp 5001")
    parser = argparse.ArgumentParser(description='manual to this script')
    parser.add_argument('-ld', type=str, default="127.0.0.1", help='local-addr 本地监听地址, 默认为: 127.0.0.1')
    parser.add_argument('-lp', type=int, default=5000, help="local-port 本地监听端口,默认为: 5000")

    parser.add_argument('-rd', type=str, help='remote-addr 远程服务端地址, 默认为: 127.0.0.1')
    parser.add_argument('-rp', type=int, default=5001, help="remote-port 远程监听端口,默认为: 5001")

    parser.add_argument('-up', type=str, default='Nothing', help="是否进行加解密操作 默认不做任何操作")
    parser.add_argument('-down', type=str, default='Nothing', help="是否进行加解密操作 默认不做任何操作")
    args = parser.parse_args()
    if not args.rd:
        print("请输入远程服务器地址")
        exit()

    t = TcpProxy(local_addr=args.ld, local_port=args.lp, remote_addr=args.rd, remote_port=args.rp, up=args.up,
                 down=args.down)
    t.run()
